`include "aes_core_h.v"

module aes_core_control_unit
(
  //OUTPUTS
  output reg  [ `COL_NUM - 1:0] col_en,      //Enable signal to column registers
  output wire [ `KEY_NUM - 1:0] key_en,      //Enable signal to key registers
  output reg config_en,                      //Enable signal to cofiguration register
  output reg config_src,                     //Select signal to cofiguration register input
  output reg key_src,                        //Select signal to key registers input
  output reg [1:0] col_src,                  //Select signal to column registers input
  output reg [2:0] sbox_src,                 //Select signal to MUX_SBOX
  output reg [1:0] rk_src,                   //Select signal to MUX_RK
  output reg [1:0] rkkey_src,                //Select signal to MUX_KEY
  output reg rd_count_clear,                 //Signal to clear round counter
  output reg rd_count_en,                    //Enable signal to round counter
  output reg start_set,                      //Input to start bit
  output reg int_set,                        //Input to interruption bit
  output wire busy_n,                        //Indicates that core is busy
  output wire key_gen,                       //Indicates that core is in key generation phase
  //output [1:0] bus_reg_src                //Select signal to MUX_BUS
  //INPUTS
  input [`COUNT_SIZE - 1:0] rd_count,    //Output of counter
  input start,                           //Start Bit
  input [1:0]mode,                       //Modes Bits
  input enc_dec,                         //Encryption/Decryption operation selector
  input clk,                             //System Clock
  input rst_n                            //Reset, asynchrounous, active low
);

//=============================================================================
// 1)  Column registers: Mux source
//=============================================================================
localparam COL_CORE = 2'b00;
localparam COL_BUS  = 2'b01;
localparam COL_SFT  = 2'b10;

//=============================================================================
// 2)  MUX_COL source
//=============================================================================
localparam MUXCOL_COL0 = 2'b00;
localparam MUXCOL_COL1 = 2'b01;
localparam MUXCOL_COL2 = 2'b10;
localparam MUXCOL_COL3 = 2'b11;

//=============================================================================
// 3)  MUX_SBOX source
//=============================================================================
localparam MUXSBOX_COL0    = 3'b000;
localparam MUXSBOX_COL1    = 3'b001;
localparam MUXSBOX_COL2    = 3'b010;
localparam MUXSBOX_COL3    = 3'b011;
localparam MUXSBOX_ROTWORD = 3'b100;

//=============================================================================
// 4)  MUX_RK source
//=============================================================================
localparam MUXRK_COL    = 2'b00;
localparam MUXRK_SBOX   = 2'b01;
localparam MUXRK_MIXCOL = 2'b10;

//=============================================================================
// 5)  KEY registers: Mux source
//=============================================================================
localparam KEY_CORE = 1'b1;
localparam KEY_BUS  = 1'b0;

//=============================================================================
// 6)  MUX_KEY source
//=============================================================================
localparam MUXKEY_KEY0 = 2'b00;
localparam MUXKEY_KEY1 = 2'b01;
localparam MUXKEY_KEY2 = 2'b10;
localparam MUXKEY_KEY3 = 2'b11;

//=============================================================================
// 7)  Start and interrupt source
//=============================================================================
localparam CONFIG_CORE = 1'b1;
localparam CONFIG_BUS  = 1'b0;

//=============================================================================
// 8)  REG_SDW_COL enable
//=============================================================================
localparam SDW_DIS  = 4'b0000;
localparam SDW_COL0 = 4'b0001;
localparam SDW_COL1 = 4'b0010;
localparam SDW_COL2 = 4'b0100;
localparam SDW_COL3 = 4'b1000;

//=============================================================================
// 9)  REG_COL enable
//=============================================================================
localparam COL_DIS = 4'b0000;
localparam COL0    = 4'b0001;
localparam COL1    = 4'b0010;
localparam COL2    = 4'b0100;
localparam COL3    = 4'b1000;
localparam COL_ALL = 4'b1111;

//=============================================================================
// 10)  REG_KEY enable
//=============================================================================
localparam KEY_DIS = 4'b0000;
localparam KEY0    = 4'b0001;
localparam KEY1    = 4'b0010;
localparam KEY2    = 4'b0100;
localparam KEY3    = 4'b1000;

//=============================================================================
// 11)  REG_CONFIG enable
//=============================================================================
localparam CONFIG_EN  = 1'b1;
localparam CONFIG_DIS = 1'b0;

//=============================================================================
// 12)  FSM STATES
//=============================================================================
localparam IDLE        = 4'd0;
localparam ROUND0_COL0 = 4'd1;
localparam ROUND0_COL1 = 4'd2;
localparam ROUND0_COL2 = 4'd3;
localparam ROUND0_COL3 = 4'd4;
localparam ROUND_KEY0  = 4'd5;
localparam ROUND_COL0  = 4'd6;
localparam ROUND_COL1  = 4'd7;
localparam ROUND_COL2  = 4'd8;
localparam ROUND_COL3  = 4'd9;
localparam READY       = 4'd10;
localparam GEN_KEY0    = 4'd11;
localparam GEN_KEY1    = 4'd12;
localparam GEN_KEY2    = 4'd13;
localparam GEN_KEY3    = 4'd14;

//=============================================================================
// 13)  Operation modes
//=============================================================================
localparam ENCRYPTION   = 2'b01;
localparam DECRYPTION   = 2'b00;
localparam KEY_EXPANDER = 2'b11;
localparam DEC_W_KEY    = 2'b10;

reg [3:0] state, next_state;
wire last_round;
reg [ `KEY_NUM - 1:0] key_en_ctu;

//State Memory
always @(posedge clk, negedge rst_n)
  begin
    if(!rst_n)
      begin
        state <= IDLE;
      end
    else
      begin
        state <= next_state;
      end
  end
  
//Next State Logic
always @(*)
  begin
    next_state = state;
    case(state)
      IDLE:
        begin
          if(!start)
            next_state = IDLE;
          else
            case(mode)
              ENCRYPTION  : next_state = ROUND0_COL0; 
              DECRYPTION  : next_state = ROUND0_COL3;
              KEY_EXPANDER: next_state = GEN_KEY0;
              DEC_W_KEY   : next_state = GEN_KEY0;
              default:      next_state = IDLE;
            endcase
        end
      ROUND0_COL0:
        begin
          next_state = enc_dec ? ROUND0_COL1 : ROUND_KEY0;
        end
      ROUND0_COL1:
        begin
          next_state = enc_dec ? ROUND0_COL2 : ROUND0_COL0;
        end
      ROUND0_COL2:
        begin
          next_state = enc_dec ? ROUND0_COL3 : ROUND0_COL1;
        end
      ROUND0_COL3:
        begin
          next_state = enc_dec ? ROUND_KEY0 : ROUND0_COL2;
        end
      ROUND_KEY0:
        begin
          if(enc_dec)
            next_state = ROUND_COL0;
          else
            next_state = (rd_count == 10) ? READY : ROUND_COL3;
        end
      ROUND_COL0:
        begin
          next_state = enc_dec ? ROUND_COL1 : ROUND_KEY0;
        end 
      ROUND_COL1:
        begin
          next_state = enc_dec ? ROUND_COL2 : ROUND_COL0;
        end 
      ROUND_COL2:
        begin
          next_state = enc_dec ? ROUND_COL3 : ROUND_COL1;
        end 
      ROUND_COL3:
        begin
          if(!enc_dec)
            next_state = ROUND_COL2;
          else
            next_state = (rd_count == 10) ? READY : ROUND_KEY0; 
        end 
      READY:
        begin
          next_state = IDLE;
        end
      GEN_KEY0:
        begin
          next_state = GEN_KEY1;
        end
      GEN_KEY1:
        begin
          next_state = GEN_KEY2;
        end
      GEN_KEY2:
        begin
          next_state = GEN_KEY3;
        end
      GEN_KEY3:
        begin
          if(rd_count == 10)
            begin
              next_state = (mode != KEY_EXPANDER) ? ROUND0_COL3 : READY;
            end
          else
            next_state = GEN_KEY0;
          //next_state = (rd_count == 10) ? READY : GEN_KEY0;
        end
      default: 
          next_state = IDLE;
    endcase
  end
  
//Output Logic 
assign busy_n = (state == IDLE);
assign last_round = (rd_count == 10);//(rd_count == 10 | (rd_count == 1 && !enc_dec)); 
assign key_gen = (state == GEN_KEY0) |
                 (state == GEN_KEY1) |
                 (state == GEN_KEY2) |
                 (state == GEN_KEY3);
                 
assign key_en = key_en_ctu | {`KEY_NUM{(start & busy_n)}};
always @(*)
  begin
    rd_count_clear = 0;
    col_src = COL_CORE;
    key_src = KEY_CORE;
    
    rk_src = MUXRK_COL;//dont care
    rkkey_src = MUXKEY_KEY0;//dont care
    
    sbox_src = MUXSBOX_COL0;
    key_en_ctu = KEY_DIS;
    col_en = COL_DIS;
    rd_count_en = 0;
    
    config_en = CONFIG_DIS;
    config_src = CONFIG_CORE;
    int_set = 0;
    start_set = 0;
    
    case(state)
      IDLE:
        begin
          rd_count_clear = 1;
          col_src = COL_BUS;
          key_src = KEY_BUS;
          config_src = CONFIG_BUS;
        end
      ROUND0_COL0:
        begin
          sbox_src = MUXSBOX_COL0;
          rk_src = MUXRK_COL;
          rkkey_src = MUXKEY_KEY0;
          col_en = COL0;
          /*if(!enc_dec)
            begin
              sbox_src = MUXSBOX_ROTWORD;
              key_src = KEY_CORE;
              key_en_ctu = KEY0;
            end*/
        end
      ROUND0_COL1:
        begin
          sbox_src = MUXSBOX_COL1;
          rk_src = MUXRK_COL;
          rkkey_src = MUXKEY_KEY1;
          col_en = COL1;
          if(!enc_dec)
            begin
              key_src = KEY_CORE;
              key_en_ctu = KEY1;
            end
        end
      ROUND0_COL2:
        begin
          sbox_src = MUXSBOX_COL2;
          rk_src = MUXRK_COL;
          rkkey_src = MUXKEY_KEY2;
          col_en = COL2;
          if(!enc_dec)
            begin
              key_src = KEY_CORE;
              key_en_ctu = KEY2;
            end
        end
      ROUND0_COL3:
        begin
          sbox_src = MUXSBOX_COL3;
          rk_src = MUXRK_COL;
          rkkey_src = MUXKEY_KEY3;
          col_en = COL3;
          if(!enc_dec)
            begin
              key_src = KEY_CORE;
              key_en_ctu = KEY3;
            end
        end
      ROUND_KEY0:
        begin
          //if(enc_dec)
            //begin
              sbox_src = MUXSBOX_ROTWORD;
              key_src = KEY_CORE;
              key_en_ctu = KEY0;
            //end
          if(enc_dec | !last_round)
            begin
              col_src = COL_SFT;
              col_en =  COL_ALL;
            end
          rd_count_en = 1;
        end
      ROUND_COL0:
        begin
          sbox_src = MUXSBOX_COL0;
          rk_src = (last_round) ? MUXRK_SBOX : MUXRK_MIXCOL;
          rkkey_src = MUXKEY_KEY0;
          col_en = COL0;
          key_src = KEY_CORE;
          if(enc_dec)
            begin
              key_en_ctu = KEY1;
            end          
        end 
      ROUND_COL1:
        begin
          sbox_src = MUXSBOX_COL1;
          rk_src = (last_round) ? MUXRK_SBOX : MUXRK_MIXCOL;
          rkkey_src = MUXKEY_KEY1;
          col_en = COL1;
          key_src = KEY_CORE;
          key_en_ctu = enc_dec ? KEY2 : KEY1;  
        end 
      ROUND_COL2:
        begin
          sbox_src = MUXSBOX_COL2;
          rk_src = (last_round) ? MUXRK_SBOX : MUXRK_MIXCOL;
          rkkey_src = MUXKEY_KEY2;
          col_en = COL2;
          key_src = KEY_CORE;
          key_en_ctu = enc_dec ? KEY3 : KEY2; 
        end 
      ROUND_COL3:
        begin
          sbox_src = MUXSBOX_COL3;
          rk_src = (last_round) ? MUXRK_SBOX : MUXRK_MIXCOL;
          rkkey_src = MUXKEY_KEY3;
          col_en = COL3;
          key_src = KEY_CORE;
          if(!enc_dec)
            begin
              key_en_ctu = KEY3;
            end 
        end 
      READY:
        begin
          config_en = CONFIG_EN;
          config_src = CONFIG_CORE;
          int_set = 1;
          start_set = 0;
        end
      GEN_KEY0:
        begin
          sbox_src = MUXSBOX_ROTWORD;
          key_src = KEY_CORE;
          key_en_ctu = KEY0;
          rd_count_en = 1;
        end
      GEN_KEY1:
        begin
          key_src = KEY_CORE;
          key_en_ctu = KEY1;
        end
      GEN_KEY2:
        begin
          key_src = KEY_CORE;
          key_en_ctu = KEY2;
        end
      GEN_KEY3:
        begin
          key_src = KEY_CORE;
          key_en_ctu = KEY3;
          if(rd_count == 10)
            rd_count_clear = 1;
        end
      default:; 
    endcase  
  end
endmodule
 